import os

import numpy as np
import torch
import torch.nn as nn
import argparse
import ast
import random
from utils import evaluate

if __name__ == '__main__':
    pictures = {}
    with open('/home/n40/code/aesthetic_path/Active_learning_Aes/aes_ability.txt', mode='r', encoding='utf-8') as file:
        lines = file.readlines()
        for line in lines:
            line = line.strip()
            name, val = line.split(' ', 1)
            val = ast.literal_eval(val)
            if name not in pictures:
                pictures[name] = {}
            pictures[name]['aes_ability'] = val
    with open('/home/n40/code/aesthetic_path/Active_learning_Aes/mark_ability.txt', mode='r', encoding='utf-8') as file:
        lines = file.readlines()
        for line in lines:
            line = line.strip()
            name, val = line.split(' ', 1)
            val = ast.literal_eval(val)
            if name not in pictures:
                pictures[name] = {}
            pictures[name]['mark_ability'] = val
    with open('/home/n40/code/aesthetic_path/Active_learning_Aes/dimension_mark.txt', mode='r', encoding='utf-8') as file:
        lines = file.readlines()
        for line in lines:
            line = line.strip()
            name, val = line.split(' ', 1)
            val = ast.literal_eval(val)
            if name not in pictures:
                pictures[name] = {}
            pictures[name]['action_list'] = val
    
    cnt, sum, tl = [0 for _ in range(77)], [0 for _ in range(77)], 0
    for ke, va in pictures.items():
        tl += 1
        for idx in va['action_list']:
            cnt[idx] += 1
            sum[idx] += abs(va['mark_ability'][idx] - va['aes_ability'][idx])
    sum = [sum[_] / cnt[_] for _ in range(len(cnt))]
    cnt = [_ / tl for _ in cnt]
    print(cnt)
    print(sum)

    mae = 0
    # while mae <= 0.153:
    # score, label = [], []
    # for ke, va in pictures.items():
    #     for idx in range(len(va['mark_ability'])):
    #         va['aes_ability'][idx] = ((va['aes_ability'][idx] - 0.2) * 1 / 0.6) + 0
    #     temp = va['aes_ability']
    #     for idx in range(len(va['mark_ability'])):
    #         va['mark_ability'][idx] = (va['mark_ability'][idx] - 0.2) / 0.15 * 0.25
    #     sum = len(va['action_list'])
    #     # excluded_set = set(va['action_list'])
    #     # all_numbers = set(range(77))
    #     # available_numbers = all_numbers - excluded_set
    #     # random_list = random.sample(available_numbers, max(1, sum // 2))
    #     random_list = random.sample(range(77), max(1, sum // 1))
    #     # print(sum, len(set(random_list)))
    #     for idx in random_list:
    #         temp[idx] = va['mark_ability'][idx]
    #     score.append(temp)
    #     label.append(va['mark_ability'])
    # score, label = torch.tensor(score), torch.tensor(label)
    # mae, mse, rmse = evaluate(score, label)
    # print(f'mae {mae} mse {mse}')
    # print(pictures)
